#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
地震响应位移和加速度时程绘图程序
绘制17、29、41节点的水平位移和加速度响应
"""

import numpy as np
import matplotlib.pyplot as plt
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

def read_response_data(filename):
    """
    读取OpenSees输出文件（位移或加速度）
    返回：时间数组，各节点响应数组
    """
    if not os.path.exists(filename):
        print(f"错误：找不到文件 {filename}")
        return None, None
    
    try:
        # 读取数据文件
        data = np.loadtxt(filename)
        
        # 第一列是时间，后续列是各节点的响应
        time = data[:, 0]
        node17_resp = data[:, 1]  # 节点17（二层）
        node29_resp = data[:, 2]  # 节点29（三层）
        node41_resp = data[:, 3]  # 节点41（四层）
        
        return time, (node17_resp, node29_resp, node41_resp)
    
    except Exception as e:
        print(f"读取文件时出错：{e}")
        return None, None

def read_earthquake_data(filename='earthquake.txt', dt=0.001):
    """
    读取地震动输入数据
    支持两种格式：
    1. 两列格式：时间 加速度
    2. 单列格式：只有加速度（需要指定时间步长dt）
    """
    if os.path.exists(filename):
        try:
            eq_data = np.loadtxt(filename)
            
            # 检查数据维度
            if eq_data.ndim == 1:
                # 单列数据，只有加速度
                eq_accel = eq_data
                eq_time = np.arange(len(eq_accel)) * dt
                return eq_time, eq_accel
            elif eq_data.shape[1] == 2:
                # 两列数据：时间 + 加速度
                eq_time = eq_data[:, 0]
                eq_accel = eq_data[:, 1]
                return eq_time, eq_accel
            elif eq_data.shape[1] == 1:
                # 单列数据（以列向量形式）
                eq_accel = eq_data[:, 0]
                eq_time = np.arange(len(eq_accel)) * dt
                return eq_time, eq_accel
            else:
                print(f"不支持的数据格式：{eq_data.shape}")
                return None, None
        except Exception as e:
            print(f"读取地震动文件时出错：{e}")
            return None, None
    return None, None

def plot_complete_earthquake_response():
    """
    绘制完整的地震响应：输入、位移、加速度
    """
    # 读取位移数据
    disp_time, disp_data = read_response_data('displacement_x.out')
    # 读取加速度数据
    accel_time, accel_data = read_response_data('acceleration_x.out')
    # 读取地震动输入（使用模型中的时间步长）
    eq_time, eq_accel = read_earthquake_data('earthquake.txt', dt=0.001)
    
    if disp_time is None:
        print("无法读取位移数据文件")
        return
    
    if accel_time is None:
        print("无法读取加速度数据文件")
        return
    
    # 解包数据
    node17_disp, node29_disp, node41_disp = disp_data
    node17_accel, node29_accel, node41_accel = accel_data
    
    # 创建三子图布局
    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(14, 12))
    
    # 子图1：地震动输入
    if eq_time is not None:
        ax1.plot(eq_time, eq_accel, 'k-', linewidth=1.2, label='地震动输入')
        ax1.set_ylabel('加速度 (m/s²)')
        ax1.set_title('地震动输入加速度时程')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
    else:
        ax1.text(0.5, 0.5, '未找到地震动输入文件', ha='center', va='center', 
                transform=ax1.transAxes, fontsize=12)
        ax1.set_title('地震动输入加速度时程')
    
    # 子图2：位移响应
    ax2.plot(disp_time, node17_disp*1000, 'b-', linewidth=1.5, label='节点17（二层）')
    ax2.plot(disp_time, node29_disp*1000, 'r-', linewidth=1.5, label='节点29（三层）')
    ax2.plot(disp_time, node41_disp*1000, 'g-', linewidth=1.5, label='节点41（四层）')
    ax2.set_ylabel('水平位移 (mm)')
    ax2.set_title('结构水平位移响应')
    ax2.grid(True, alpha=0.3)
    ax2.legend()
    
    # 子图3：加速度响应
    ax3.plot(accel_time, node17_accel, 'b-', linewidth=1.5, label='节点17（二层）')
    ax3.plot(accel_time, node29_accel, 'r-', linewidth=1.5, label='节点29（三层）')
    ax3.plot(accel_time, node41_accel, 'g-', linewidth=1.5, label='节点41（四层）')
    ax3.set_xlabel('时间 (s)')
    ax3.set_ylabel('水平加速度 (m/s²)')
    ax3.set_title('结构水平加速度响应')
    ax3.grid(True, alpha=0.3)
    ax3.legend()
    
    plt.tight_layout()
    plt.savefig('earthquake_complete_response.png', dpi=300, bbox_inches='tight')
    print("完整响应图已保存为 earthquake_complete_response.png")
    plt.show()
    
    # 输出统计信息
    print_statistics(disp_time, disp_data, accel_time, accel_data)

def plot_response_comparison():
    """
    绘制各楼层响应对比图
    """
    # 读取数据
    disp_time, disp_data = read_response_data('displacement_x.out')
    accel_time, accel_data = read_response_data('acceleration_x.out')
    
    if disp_time is None or accel_time is None:
        print("无法读取响应数据文件")
        return
    
    # 解包数据
    node17_disp, node29_disp, node41_disp = disp_data
    node17_accel, node29_accel, node41_accel = accel_data
    
    # 创建2x2子图布局
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
    
    # 位移时程对比
    ax1.plot(disp_time, node17_disp*1000, 'b-', linewidth=1.5, label='节点17（二层）')
    ax1.plot(disp_time, node29_disp*1000, 'r-', linewidth=1.5, label='节点29（三层）')
    ax1.plot(disp_time, node41_disp*1000, 'g-', linewidth=1.5, label='节点41（四层）')
    ax1.set_xlabel('时间 (s)')
    ax1.set_ylabel('水平位移 (mm)')
    ax1.set_title('水平位移时程对比')
    ax1.grid(True, alpha=0.3)
    ax1.legend()
    
    # 加速度时程对比
    ax2.plot(accel_time, node17_accel, 'b-', linewidth=1.5, label='节点17（二层）')
    ax2.plot(accel_time, node29_accel, 'r-', linewidth=1.5, label='节点29（三层）')
    ax2.plot(accel_time, node41_accel, 'g-', linewidth=1.5, label='节点41（四层）')
    ax2.set_xlabel('时间 (s)')
    ax2.set_ylabel('水平加速度 (m/s²)')
    ax2.set_title('水平加速度时程对比')
    ax2.grid(True, alpha=0.3)
    ax2.legend()
    
    # 最大位移对比
    max_disps = [np.max(np.abs(node17_disp)), 
                 np.max(np.abs(node29_disp)), 
                 np.max(np.abs(node41_disp))]
    floors = ['二层\n(节点17)', '三层\n(节点29)', '四层\n(节点41)']
    
    bars1 = ax3.bar(floors, np.array(max_disps)*1000, 
                    color=['blue', 'red', 'green'], alpha=0.7)
    ax3.set_ylabel('最大水平位移 (mm)')
    ax3.set_title('各楼层最大水平位移对比')
    ax3.grid(True, alpha=0.3, axis='y')
    
    # 在柱状图上添加数值标签
    for bar, disp in zip(bars1, max_disps):
        height = bar.get_height()
        ax3.text(bar.get_x() + bar.get_width()/2., height + height*0.02,
                f'{disp*1000:.3f}', ha='center', va='bottom')
    
    # 最大加速度对比
    max_accels = [np.max(np.abs(node17_accel)), 
                  np.max(np.abs(node29_accel)), 
                  np.max(np.abs(node41_accel))]
    
    bars2 = ax4.bar(floors, max_accels, 
                    color=['blue', 'red', 'green'], alpha=0.7)
    ax4.set_ylabel('最大水平加速度 (m/s²)')
    ax4.set_title('各楼层最大水平加速度对比')
    ax4.grid(True, alpha=0.3, axis='y')
    
    # 在柱状图上添加数值标签
    for bar, accel in zip(bars2, max_accels):
        height = bar.get_height()
        ax4.text(bar.get_x() + bar.get_width()/2., height + height*0.02,
                f'{accel:.3f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('earthquake_response_comparison.png', dpi=300, bbox_inches='tight')
    print("响应对比图已保存为 earthquake_response_comparison.png")
    plt.show()

def print_statistics(disp_time, disp_data, accel_time, accel_data):
    """
    输出统计信息
    """
    node17_disp, node29_disp, node41_disp = disp_data
    node17_accel, node29_accel, node41_accel = accel_data
    
    print("\n" + "="*50)
    print("地震响应统计结果")
    print("="*50)
    print(f"分析时长：{disp_time[-1]:.2f} 秒")
    print(f"时间步数：{len(disp_time)} 步")
    
    print("\n位移响应统计：")
    print(f"节点17（二层）最大位移：{np.max(np.abs(node17_disp))*1000:.3f} mm")
    print(f"节点29（三层）最大位移：{np.max(np.abs(node29_disp))*1000:.3f} mm")
    print(f"节点41（四层）最大位移：{np.max(np.abs(node41_disp))*1000:.3f} mm")
    
    print("\n加速度响应统计：")
    print(f"节点17（二层）最大加速度：{np.max(np.abs(node17_accel)):.3f} m/s²")
    print(f"节点29（三层）最大加速度：{np.max(np.abs(node29_accel)):.3f} m/s²")
    print(f"节点41（四层）最大加速度：{np.max(np.abs(node41_accel)):.3f} m/s²")
    
    # 计算放大系数（如果有地震动输入）
    eq_time, eq_accel = read_earthquake_data('earthquake.txt', dt=0.001)
    if eq_accel is not None:
        max_input_accel = np.max(np.abs(eq_accel))
        print(f"\n地震动输入最大加速度：{max_input_accel:.3f} m/s²")
        print("\n加速度放大系数：")
        print(f"节点17（二层）：{np.max(np.abs(node17_accel))/max_input_accel:.2f}")
        print(f"节点29（三层）：{np.max(np.abs(node29_accel))/max_input_accel:.2f}")
        print(f"节点41（四层）：{np.max(np.abs(node41_accel))/max_input_accel:.2f}")

def main():
    """
    主函数
    """
    print("地震响应分析绘图程序（位移+加速度）")
    print("=" * 50)
    
    # 检查当前目录
    current_dir = os.getcwd()
    print(f"当前工作目录：{current_dir}")
    
    # 检查文件是否存在
    files_to_check = ['displacement_x.out', 'acceleration_x.out', 'earthquake.txt']
    for filename in files_to_check:
        if os.path.exists(filename):
            print(f"✓ 找到文件：{filename}")
        else:
            print(f"✗ 缺少文件：{filename}")
    
    print("\n开始绘制响应图...")
    
    # 绘制完整响应图
    plot_complete_earthquake_response()
    
    # 绘制响应对比图
    plot_response_comparison()
    
    print("\n绘图完成！")

if __name__ == "__main__":
    main()